Expose weights_only argument for loading checkpoints in Fabric.load() and Fabric.load_raw()
#21470
+80
−10
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
CodeVishal-17
requested review from
ethanwharris,
justusschock,
lantiga and
tchaton
as code owners
for more information, see https://pre-commit.ci
GdoongMathew
suggested changes
Jan 6, 2026
…hal-17/pytorch-lightning into fabric-load-weights-only
for more information, see https://pre-commit.ci
GdoongMathew
reviewed
Jan 6, 2026
…/CodeVishal-17/pytorch-lightning into fabric-load-weights-only" This reverts commit 6143e65, reversing changes made to debde18.
…hal-17/pytorch-lightning into fabric-load-weights-only
for more information, see https://pre-commit.ci
GdoongMathew
approved these changes
Jan 7, 2026
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #21470 +/- ##
=======================================
- Coverage 87% 87% -0%
=======================================
Files 270 270
Lines 24059 24059
=======================================
- Hits 20858 20855 -3
- Misses 3201 3204 +3 |
bhimrazy
changed the title
weights_only argument for loading checkpoints in Fabric.load() and Fabric.load_raw()
bhimrazy
approved these changes
Jan 12, 2026
tchaton
approved these changes
Jan 12, 2026
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Exposes the
weights_onlyargument inFabric.load()and forwards it to the underlying strategy checkpoint loading call.This restores compatibility with mixed-object checkpoints after the PyTorch 2.6 change where
torch.loaddefaults toweights_only=True.Motivation and context
With PyTorch 2.6,
torch.loaddefaults toweights_only=True, which prevents loading checkpoints that contain non-tensor objects (such as optimizers, schedulers, or metadata).While
strategy.load_checkpoint(..., weights_only=False)already supports this use case,Fabric.load()did not expose the argument, making it impossible for users to opt out of the new default behavior when using Fabric.This PR adds a keyword-only
weights_onlyargument toFabric.load()without changing the existing default behavior.Fixes #21459
Before submitting
PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:
Reviewer checklist
📚 Documentation preview 📚: https://pytorch-lightning--21470.org.readthedocs.build/en/21470/